"""
Qwen Agent + SiliconFlow API 集成示例
参考官方示例，使用 Assistant 类实现工具调用和自然对话
"""

import json
import requests
from qwen_agent.tools import BaseTool
from qwen_agent.agents import Assistant
from qwen_agent.utils.output_beautify import typewriter_print


# ==================== SiliconFlow API 配置 ====================

SILICONFLOW_CONFIG = {
    "api_key": "sk-ironfacdlepjcceslbwpnugxwvniqlfoluxjwnnkrfuufnef",
    "base_url": "https://api.siliconflow.cn/v1",
    "model": "Qwen/Qwen2.5-7B-Instruct"  # 可以根据需要更换模型
}


# ==================== 工具定义 ====================

class GetWeather(BaseTool):
    """天气查询工具类"""
    
    # 必须设置的类属性
    name = "get_weather"
    description = "根据城市名称获取当前天气信息"
    parameters = {
        "type": "object",
        "properties": {
            "city": {
                "type": "string",
                "description": "城市的名称，例如'北京'、'上海'等"
            }
        },
        "required": ["city"]
    }
    
    def call(self, params: dict, **kwargs) -> str:
        """
        根据指定的城市名称获取实时天气数据。
        
        Args:
            params: 包含城市名称的参数字典，格式为 {"city": "城市名"}
            **kwargs: 其他关键字参数（Assistant 可能会传递 messages 等）
            
        Returns:
            str: 包含天气信息的 JSON 字符串。
        """
        # 解析参数
        if isinstance(params, str):
            params = json.loads(params)
        
        city = params.get("city", "")
        
        # 实际的工具逻辑，这里使用模拟数据
        weather_map = {
            "北京": {"city": "北京", "temperature": "25°C", "condition": "晴", "humidity": "60%"},
            "上海": {"city": "上海", "temperature": "28°C", "condition": "多云", "humidity": "65%"},
            "深圳": {"city": "深圳", "temperature": "30°C", "condition": "晴", "humidity": "70%"},
        }
        
        weather_data = weather_map.get(
            city, 
            {"city": city, "temperature": "未知", "condition": "未知", "humidity": "未知"}
        )
        
        return json.dumps(weather_data, ensure_ascii=False)


# ==================== Agent 初始化 ====================

def init_agent_service():
    """
    初始化 Agent 服务，使用 SiliconFlow API 作为 LLM 后端
    
    Returns:
        Assistant: 配置好的 Assistant 实例
    """
    # LLM 配置：使用 OpenAI 兼容的 API（SiliconFlow）
    llm_cfg = {
        'model': SILICONFLOW_CONFIG['model'],
        'model_server': SILICONFLOW_CONFIG['base_url'],
        'api_key': SILICONFLOW_CONFIG['api_key'],
    }
    
    # 工具列表
    tools = [
        GetWeather(),  # 自定义天气查询工具
        # 可以添加更多工具，例如：
        # 'code_interpreter',  # 内置代码解释器
    ]
    
    # 创建 Assistant
    bot = Assistant(
        llm=llm_cfg,
        function_list=tools,
        name='Qwen Agent + SiliconFlow 示例',
        description="我是一个智能助手，可以帮助你查询天气信息，也可以进行自然对话。"
    )
    
    return bot


# ==================== 测试函数 ====================

def test(query: str = '查询北京的天气'):
    """
    测试函数：单次对话测试
    
    Args:
        query: 用户查询内容
    """
    bot = init_agent_service()
    
    messages = [{'role': 'user', 'content': query}]
    response_plain_text = ''
    
    print(f"用户: {query}")
    print("助手: ", end="", flush=True)
    
    for response in bot.run(messages=messages):
        response_plain_text = typewriter_print(response, response_plain_text)
    
    print()  # 换行


def test_siliconflow_api():
    """测试 SiliconFlow API 直接调用"""
    url = f"{SILICONFLOW_CONFIG['base_url']}/chat/completions"
    
    payload = {
        "model": SILICONFLOW_CONFIG["model"],
        "messages": [
            {
                "role": "user",
                "content": "你好，请简单介绍一下你自己。"
            }
        ]
    }
    
    headers = {
        "Authorization": f"Bearer {SILICONFLOW_CONFIG['api_key']}",
        "Content-Type": "application/json"
    }
    
    try:
        response = requests.post(url, json=payload, headers=headers)
        response.raise_for_status()
        result = response.json()
        print("API 响应:", json.dumps(result, ensure_ascii=False, indent=2))
        return result
    except requests.exceptions.RequestException as e:
        print(f"API 调用失败: {e}")
        return None


def test_weather_tool():
    """测试天气查询工具"""
    weather_tool = GetWeather()
    
    # 测试不同城市的天气查询
    test_cities = ["北京", "上海", "深圳"]
    
    print("\n=== 天气查询工具测试 ===")
    for city in test_cities:
        params = {"city": city}
        result = weather_tool.call(params)
        print(f"{city}: {result}")


def app_tui():
    """
    终端交互式对话（TUI）
    参考官方示例的 app_tui 函数
    """
    bot = init_agent_service()
    
    print("\n" + "=" * 60)
    print("进入交互式对话模式")
    print("输入 'quit' 或 'exit' 退出")
    print("=" * 60)
    
    messages = []
    while True:
        try:
            query = input('\n你: ').strip()
            
            if not query:
                continue
            
            if query.lower() in ['quit', 'exit', '退出']:
                print("\n再见！")
                break
            
            messages.append({'role': 'user', 'content': query})
            
            response = []
            response_plain_text = ''
            
            print("助手: ", end="", flush=True)
            for response in bot.run(messages=messages):
                response_plain_text = typewriter_print(response, response_plain_text)
            
            print()  # 换行
            
            # 更新消息历史
            messages.extend(response)
            
            # 限制历史长度
            if len(messages) > 20:
                messages = messages[-20:]
                
        except KeyboardInterrupt:
            print("\n\n再见！")
            break
        except Exception as e:
            print(f"\n错误: {e}")
            import traceback
            traceback.print_exc()


def test_conversation():
    """测试对话功能（工具使用 + 自然对话）"""
    print("\n" + "=" * 60)
    print("Agent 对话测试（工具使用 + 自然对话）")
    print("=" * 60)
    
    bot = init_agent_service()
    messages = []
    
    # 混合对话场景：既有工具使用，也有自然对话
    test_scenarios = [
        {
            "type": "自然对话",
            "message": "你好，请介绍一下你自己"
        },
        {
            "type": "工具使用",
            "message": "北京的天气怎么样？"
        },
        {
            "type": "自然对话",
            "message": "今天天气不错，你觉得适合做什么活动？"
        },
        {
            "type": "工具使用",
            "message": "帮我查询一下上海的天气"
        },
        {
            "type": "自然对话",
            "message": "如果下雨的话，我应该带什么？"
        },
        {
            "type": "工具使用",
            "message": "深圳的天气如何？"
        },
    ]
    
    for i, scenario in enumerate(test_scenarios, 1):
        print(f"\n[场景 {i}] {scenario['type']}")
        print(f"用户: {scenario['message']}")
        print("-" * 60)
        
        try:
            messages.append({'role': 'user', 'content': scenario['message']})
            
            response_plain_text = ''
            print("助手: ", end="", flush=True)
            
            for response in bot.run(messages=messages):
                response_plain_text = typewriter_print(response, response_plain_text)
            
            print()  # 换行
            
            # 更新消息历史
            messages.extend(response)
            
            # 限制历史长度
            if len(messages) > 20:
                messages = messages[-20:]
                
        except Exception as e:
            print(f"错误: {e}")
            import traceback
            traceback.print_exc()
    
    print("\n" + "=" * 60)
    print("对话测试完成！")
    print("=" * 60)


# ==================== 主函数 ====================

def main():
    """主函数"""
    import sys
    
    print("=" * 60)
    print("Qwen Agent + SiliconFlow API 集成示例")
    print("=" * 60)
    
    # 如果提供了命令行参数，运行对应的测试
    if len(sys.argv) > 1:
        mode = sys.argv[1].lower()
        
        if mode == "interactive" or mode == "chat" or mode == "tui":
            # 交互式对话模式
            app_tui()
        elif mode == "conversation" or mode == "test":
            # 预设对话场景测试
            test_conversation()
        elif mode == "tool":
            # 测试工具
            test_weather_tool()
        elif mode == "api":
            # 测试 API
            test_siliconflow_api()
        elif mode == "single":
            # 单次测试
            test("查询北京的天气")
        else:
            print(f"未知模式: {mode}")
            print("\n可用模式:")
            print("  python tool_ues_test.py interactive  - 交互式对话（TUI）")
            print("  python tool_ues_test.py conversation - 预设对话场景")
            print("  python tool_ues_test.py tool        - 测试工具")
            print("  python tool_ues_test.py api         - 测试 API")
            print("  python tool_ues_test.py single     - 单次对话测试")
    else:
        # 默认运行：先测试工具，再运行对话测试
        print("\n1. 测试天气工具...")
        test_weather_tool()
        
        print("\n2. 运行对话测试（工具使用 + 自然对话）...")
        test_conversation()
        
        print("\n" + "=" * 60)
        print("程序执行完成！")
        print("=" * 60)
        print("\n提示：")
        print("  - 运行 'python tool_ues_test.py interactive' 进入交互式对话模式")
        print("  - 运行 'python tool_ues_test.py conversation' 查看预设对话场景")
        print("  - 运行 'python tool_ues_test.py single' 进行单次对话测试")


if __name__ == "__main__":
    main()
